import numpy as np

from src.impl.output import get_output_float16, get_output_float32


def has_precision_problem_occurred(tensor, operator, pattern='MRE', b=0.0001, r=0.1, epsilon= 1e-7):
    # 得到tf和pytorch输出的numpy数组
    tf_output_b, torch_output_b, mnn_output_b, variable = get_output_float32(tensor, operator)

    # 获得三组输出的差异
    tf_torch_abs = np.maximum(tf_output_b - torch_output_b, torch_output_b - tf_output_b)
    tf_mnn_abs = np.maximum(tf_output_b - mnn_output_b, mnn_output_b - tf_output_b)
    torch_mnn_abs = np.maximum(torch_output_b - mnn_output_b, mnn_output_b - torch_output_b)
    # 获得总和
    tf_torch_var = np.sum(tf_torch_abs)
    tf_mnn_var = np.sum(tf_mnn_abs)
    torch_mnn_var = np.sum(torch_mnn_abs)
    # 比较并得到benchmark
    if tf_torch_var <= tf_mnn_var and tf_torch_var <= torch_mnn_var:
        benchmark = tf_output_b / 2 + torch_output_b / 2
    elif tf_mnn_var <= tf_torch_var and tf_mnn_var <= torch_mnn_var:
        benchmark = tf_output_b / 2 + mnn_output_b / 2
    else:
        benchmark = torch_output_b / 2 + mnn_output_b / 2
    tf_output_16, torch_output_16, mnn_output_16, _ = get_output_float16(tensor, operator, variable)

    # 差异矩阵
    diff_matrix_tf = tf_output_16 - benchmark
    diff_matrix_torch = torch_output_16 - benchmark
    diff_matrix_mnn = mnn_output_16 - benchmark

    # 求和项(绝对值)
    # epsilon取一个很小的值防止得0
    single_item_matrix_abs_tf = np.maximum(diff_matrix_tf, -diff_matrix_tf) / \
                            (np.maximum(tf_output_16, -tf_output_16) + epsilon)
    single_item_matrix_abs_torch = np.maximum(diff_matrix_torch, -diff_matrix_torch) / \
                            (np.maximum(torch_output_16, -torch_output_16) + epsilon)
    single_item_matrix_abs_mnn = np.maximum(diff_matrix_mnn, -diff_matrix_mnn) / \
                            (np.maximum(mnn_output_16, -mnn_output_16) + epsilon)

    # 算MRE或MARE
    if pattern == 'MRE':
        # print("tf: {0}".format(np.sum(single_item_matrix_abs_tf) / np.size(single_item_matrix_abs_tf)))
        # print("torch: {0}".format(np.sum(single_item_matrix_abs_torch) / np.size(single_item_matrix_abs_torch)))
        # print(np.sum(single_item_matrix_abs_tf) / np.size(single_item_matrix_abs_tf) > b)
        # print(np.sum(single_item_matrix_abs_tf) / np.size(single_item_matrix_abs_tf))
        return np.sum(single_item_matrix_abs_tf) / np.size(single_item_matrix_abs_tf) > b, \
               np.sum(single_item_matrix_abs_torch) / np.size(single_item_matrix_abs_torch) > b, \
               np.sum(single_item_matrix_abs_mnn) / np.size(single_item_matrix_abs_mnn) > b
    elif pattern == 'MARE':
        # print("tf: {0}".format(np.max(single_item_matrix_abs_tf)))
        # print("torch: {0}".format(np.max(single_item_matrix_abs_torch)))
        # # print(np.max(single_item_matrix_abs_tf))
        # print(np.max(single_item_matrix_abs_torch))
        return np.max(single_item_matrix_abs_tf) > r, np.max(single_item_matrix_abs_torch) > r, \
               np.max(single_item_matrix_abs_mnn) > r